﻿using System.Collections.Generic;
using System.Linq;
using UnityEngine;

namespace CWM.Skinn
{
    public static partial class SkinnEx
    {

        public static bool GetMeshBoneIslands(SkinnedMeshRenderer skinnedMesh, out Mesh[] meshes, bool[] bones = null)
        {
            meshes = new Mesh[2];
            if (IsNullOrNotVailid(skinnedMesh)) return false;

            if(IsNullOrEmpty(bones))
            {
                Mesh[] allBones;
                if (GetMeshBoneIslands(skinnedMesh, out allBones)) { meshes = allBones; return true; }
                else return false;
            }

            var source = skinnedMesh.sharedMesh;
            int[] boneIndices;
            if (!BoneIslandUtility.GetBoneIndices(source, out boneIndices, 16)) return false;
            var triangles = source.triangles;
            var tempMesh = source.Clone() as Mesh;
            tempMesh.subMeshCount = 1;
            var tempHashes = new List<Vector3Int>(source.triangles.Length / 3);
            var indices = new List<int>(triangles.Length);
            for (int i = 0; i < bones.Length; i++)
            {
                if (!bones[i]) continue;
                for (int ii = 0; ii < triangles.Length; ii += 3)
                {
                    int a = triangles[ii + 0], b = triangles[ii + 1], c = triangles[ii + 2];
                    var count = boneIndices[a] == i ? 1 : 0;
                    count += boneIndices[b] == i ? 1 : 0;
                    count += boneIndices[c] == i ? 1 : 0;
                    if (count > 0)
                    {
                        var tris = new Vector3Int(a, b, c);
                        if (!tempHashes.Contains(tris)) { tempHashes.Add(tris); indices.Add(a); indices.Add(b); indices.Add(c); };
                    }
                }
            }

            if (indices.Count > 2)
            {
                tempMesh.SetTriangles(indices.ToArray(), 0);
                meshes[0] = tempMesh.GetSubmesh(0);
            }

            indices.Clear();
            for (int ii = 0; ii < triangles.Length; ii += 3)
            {
                int a = triangles[ii + 0], b = triangles[ii + 1], c = triangles[ii + 2];
                var tris = new Vector3Int(a, b, c);
                if (!tempHashes.Contains(tris)) { tempHashes.Add(tris); indices.Add(a); indices.Add(b); indices.Add(c); };
            }

            if (indices.Count > 2)
            {
                tempMesh.SetTriangles(indices.ToArray(), 0);
                meshes[1] = tempMesh.GetSubmesh(0);
            }

            Release(tempMesh);
            return meshes[0];
        }

        public static bool GetMeshBoneIslands(SkinnedMeshRenderer skinnedMesh, out Mesh[] meshes)
        {
            meshes = null;
            if (IsNullOrNotVailid(skinnedMesh)) return false;

            var source = skinnedMesh.sharedMesh;
            int[] boneIndices;
            if (!BoneIslandUtility.GetBoneIndices(source, out boneIndices, 64)) return false;
            var subMeshCount = skinnedMesh.bones.Length;
            var triangles = source.triangles;
            var tempMesh = source.Clone() as Mesh;
            tempMesh.subMeshCount = 1;
            var newMeshes = new List<Mesh>(source.bindposes.Length);
            var tempHashes = new List<Vector3Int>(source.triangles.Length / 3);

            for (int i = 0; i < subMeshCount; i++)
            {
                var indices = new List<int>();
                for (int ii = 0; ii < triangles.Length; ii += 3)
                {
                    int a = triangles[ii + 0], b = triangles[ii + 1], c = triangles[ii + 2];
                    if (boneIndices[a] == i || boneIndices[b] == i || boneIndices[b] == i)
                    {
                        var tris = new Vector3Int(a, b, c);
                        if (!tempHashes.Contains(tris))
                        {
                            tempHashes.Add(tris); indices.Add(a); indices.Add(b); indices.Add(c);
                        }
                    }
                }

                if(indices.Count > 2)
                {
                    tempMesh.SetTriangles(indices.ToArray(), 0);
                    var boneMesh = tempMesh.GetSubmesh(0);
                    boneMesh.name = EnforceObjectName(skinnedMesh.bones[i]);
                    newMeshes.Add(boneMesh);
                }
            }
            meshes = newMeshes.ToArray();
            Release(tempMesh);
            return meshes.Length > 0;

        }

        internal static class BoneIslandUtility
        {
            public static bool GetBoneIndices(Mesh source, out int[] boneIndices, int adjacentCount = 16)
            {
                boneIndices = null;

                if (!source || IsNullOrEmpty(source.boneWeights, source.vertices.Length) || IsNullOrEmpty(source.bindposes)) return false;

                var vertexCount = source.vertices.Length;
                var boneCount = source.bindposes.Length;
                adjacentCount = Mathf.Clamp(adjacentCount, 3, 64);
                int[,] adjacentIndices;
                if (!source.GetAdjacentIndices(out adjacentIndices, adjacentCount)) return false;

                var boneWeights = source.boneWeights;
                boneIndices = new int[vertexCount];
                var triangles = source.triangles;
                var weights = new float[vertexCount, boneCount];

                for (int i = 0; i < vertexCount; i ++)
                {
                    for (int iii = 0; iii < 4; iii++)
                    {
                        float w = boneWeights[i].GetBoneWeight(iii);
                        if (w > 0.245f) weights[i, boneWeights[i].GetBoneIndex(iii)] += w;
                    }

                    var count = 0f;
                    for (int ii = 0; ii < adjacentCount; ii++)
                    {
                        int a = adjacentIndices[i, ii];
                        for (int iii = 0; iii < 3; iii++)
                        {
                            if (a < 0) continue;
                            float w = boneWeights[a].GetBoneWeight(iii);
                            if (w > 0.245f) count++;
                        }
                    }

                    for (int ii = 0; ii < adjacentCount; ii++)
                    {
                        int a = adjacentIndices[i, ii];
                        if (a < 0) continue;

                        for (int iii = 0; iii < 4; iii++)
                        {
                            float w = boneWeights[a].GetBoneWeight(iii);
                            if (w > 0.245f) weights[a, boneWeights[a].GetBoneIndex(iii)] += w / count;
                        }
                    }
                }

                for (int i = 0; i < triangles.Length; i += 3)
                {
                    int a = triangles[i + 0], b = triangles[i + 1], c = triangles[i + 2];

                    var maxWeight = 0f;
                    var maxIndex = 0;

                    for (int ii = 0; ii < boneCount; ii++)
                    {
                        if (weights[a, ii] > maxWeight) { maxWeight = weights[a, ii]; maxIndex = ii; }
                        if (weights[b, ii] > maxWeight) { maxWeight = weights[b, ii]; maxIndex = ii; }
                        if (weights[c, ii] > maxWeight) { maxWeight = weights[c, ii]; maxIndex = ii; }
                    }

                    boneIndices[triangles[i + 0]] = maxIndex;
                    boneIndices[triangles[i + 1]] = maxIndex;
                    boneIndices[triangles[i + 2]] = maxIndex;
                }
                return true;
            }
        }
    }
}